import random

import torch
from matplotlib import pyplot as plt
from torchvision.datasets import MNIST, FashionMNIST
from torch.utils.data import Subset
from torchvision import transforms as T

# 这种数据处理的方式常用于
# 1.快速测试模型性能，加上你好数据量以加速训练
# 2.创建简化版数据集用于快速演示
# 3.当某种类别严格不能过多导致类别不平衡的时候

tf1 = T.Compose([
    T.Resize((32, 32), antialias=True),
    T.ToTensor()
])


def tf2(label):
    return torch.tensor(label)


# 1.加载数据集
ds = FashionMNIST("./xiazai",
           download=False,
           train=True,
           transform=tf1,
           target_transform=tf2
           )

ds_loader = torch.utils.data.DataLoader(ds, batch_size=64, shuffle=True)

for img, label in ds_loader:
    print(f"图像的大小 {img.shape} 对应的标签：{label.shape}")
    for i in range(6):
        plt.subplot(2, 3, i+1)
        plt.imshow(img[i].permute(1, 2, 0), cmap='gray')
        plt.title(label[i].item())
        plt.axis('off')
    break
plt.show()